import os
import cv2

class FrameExtractor:
    """
    视频帧提取工具类，支持单个视频或视频文件夹提取。
    新增“提取模式”参数：按秒（time）或按帧（frame）提取。
    """
    def __init__(self, input_path="", output_folder="", interval=0.05, mode="time"):
        self.input_path = input_path
        self.output_folder = output_folder or os.path.join(os.path.expanduser("~"), "Desktop", "Extracted_Frames")
        # 模式：'time' 表示按秒提取，'frame' 表示按帧提取
        if mode not in ("time", "frame"):
            raise ValueError("模式必须是 'time'（按秒）或 'frame'（按帧）")
        self.mode = mode
        # 根据模式解释 interval：time 模式下为秒数，frame 模式下为帧数
        if self.mode == "frame":
            self.interval = int(interval)
            if self.interval <= 0:
                raise ValueError("帧间隔必须大于 0")
        else:
            self.interval = interval
            if self.interval <= 0:
                raise ValueError("秒间隔必须大于 0")
        self.video_exts = (".mp4", ".avi", ".mkv", ".mov", ".flv", ".ts")

        # 自动创建输出目录
        os.makedirs(self.output_folder, exist_ok=True)

    def set_input_path(self, path: str):
        self.input_path = path
        return self

    def set_output_folder(self, folder: str):
        self.output_folder = folder
        os.makedirs(self.output_folder, exist_ok=True)
        return self

    def set_interval(self, seconds: float):
        """设置按秒提取模式及间隔（秒）。"""
        if seconds <= 0:
            raise ValueError("秒间隔必须大于 0")
        self.mode = "time"
        self.interval = seconds
        return self
    def set_mode(self, mode: str):
        if mode not in ("time", "frame"):
            raise ValueError("模式必须是 'time'（按秒）或 'frame'（按帧）")
        """设置按秒提取模式及间隔（秒）。"""
        self.mode = mode
        return self

    def set_frame_interval(self, frames: int):
        """设置按帧提取模式及帧间隔（帧数）。"""
        if frames <= 0:
            raise ValueError("帧间隔必须大于 0")
        self.mode = "frame"
        self.interval = int(frames)
        return self

    def extract(self):
        if not self.input_path or not os.path.exists(self.input_path):
            print("❌ 输入路径无效，请先设置正确的视频文件或文件夹路径。")
            return

        # 收集所有视频文件
        if os.path.isfile(self.input_path):
            video_files = [self.input_path]
        else:
            video_files = [
                os.path.join(self.input_path, f)
                for f in os.listdir(self.input_path)
                if f.lower().endswith(self.video_exts)
            ]

        if not video_files:
            print("⚠️ 未找到任何视频文件。")
            return

        for video_file in video_files:
            print(f"🎬 正在处理：{os.path.basename(video_file)}")
            self._extract_from_video(video_file)

        print(f"\n🎉 所有视频提取完成！\n📁 输出目录：{self.output_folder}")
        return self.output_folder

    def _extract_from_video(self, video_path: str):
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f"❌ 无法打开视频文件：{video_path}")
            return

        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        duration = total_frames / fps if fps > 0 else 0.0

        # 为当前视频创建输出目录
        video_name = os.path.splitext(os.path.basename(video_path))[0]
        video_output = os.path.join(self.output_folder, video_name)
        os.makedirs(video_output, exist_ok=True)

        frame_id = 0
        saved_count = 0

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            if self.mode == "time":
                # 计算当前帧对应的时间（秒）
                current_time = frame_id / fps if fps > 0 else 0
                # 判断是否达到下一个提取时间点
                if current_time >= saved_count * self.interval:
                    # 保存帧
                    output_file = os.path.join(video_output, f"frame_{saved_count:05d}.png")
                    cv2.imwrite(output_file, frame)
                    saved_count += 1

            elif self.mode == "frame":
                # 如果当前帧号满足间隔条件，则保存
                if frame_id % self.interval == 0:
                    output_file = os.path.join(video_output, f"frame_{saved_count:05d}.png")
                    cv2.imwrite(output_file, frame)
                    saved_count += 1

            frame_id += 1

        cap.release()
        print(f"✅ {os.path.basename(video_path)} 提取完成，共 {saved_count} 帧")
if __name__ == "__main__":
    extractor = FrameExtractor( input_path=r"C:\Users\33720\Desktop\temp\divide\1.mp4",output_folder =r"C:\Users\33720\Desktop\temp\divide\3")
    extractor.set_mode("time")
    extractor.set_interval(0.0)
    extractor.extract()